import torch
import torch.nn as nn
import math


class AbsolutePositionalEncoding(nn.Module):
    """绝对位置编码：显示引入位置信息，以便模型能够区分序列中的不同位置。
    思想是在每个输入序列的元素上添加一个位置向量，以表示该元素在序列中的具体位置。位置向量使用正余弦函数，生成的编码具有较强周期性，能够捕捉序列中的相对位置信息。"""
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        # 创建一个形状为（max_len, d_model）的零张量pe，用于存储所有可能的位置编码。
        pe = torch.zeros(max_len, d_model)
        # 生成一个从0到max_len-1的浮点数序列，表示序列中的每个位置，unsqueeze(1)将其形状从[max_len]变为[max_len, 1]，方便后续计算进行广播
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        # torch.arange(0, d_model, 2).float() 生成一个步长为 2 的索引序列，如 [0, 2, 4, ..., d_model-2]，表示偶数索引的维度。
        # math.log(10000.0) / d_model 计算缩放因子，使得不同维度的频率分布在不同的范围内。
        # torch.exp(...) 计算指数，以确保较低维度的正余弦波动较快，高维度的波动较慢。
        # div_term 控制了不同维度的波长，使得编码既可以捕捉局部信息，也可以捕捉全局信息。
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        # pe[:, 0::2] = torch.sin(position * div_term)：
        # 计算偶数索引维度（0, 2, 4, ...）的正弦编码。
        # position * div_term 计算不同位置和不同维度上的角度。
        # torch.sin(...) 计算正弦值。
        # pe[:, 1::2] = torch.cos(position * div_term)：
        # 计算奇数索引维度（1, 3, 5, ...）的余弦编码。
        # torch.cos(...) 计算余弦值。
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # unsqueeze(0) 在第 0 维添加一个批次维度，使得 pe 的形状从 [max_len, d_model] 变为 [1, max_len, d_model]，方便与输入 x 进行广播。
        pe = pe.unsqueeze(0)
        # register_buffer 告诉 PyTorch 该张量 pe 不是可训练参数（不会被 optimizer 更新），但仍然属于模型的一部分，并可以随模型一起保存和加载。
        self.register_buffer('pe', pe)

    def forward(self, x):
        # self.pe[:, :x.size(1)] 选取 pe 的前 seq_len 个位置，以匹配 x 的序列长度。
        # .detach() 确保 pe 在计算过程中不会参与梯度计算，避免对 pe 进行反向传播更新。
        # x + self.pe[...] 直接将位置编码加到输入 x 上。
        return x + self.pe[:, :x.size(1)].detach()


class RelativePositionalEncoding(nn.Module):
    """相对位置编码：为序列中每个位置计算相对位置差值，并使用编码函数（正余弦）编码"""
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        self.d_model = d_model
        self.max_len = max_len

        # 生成相对位置索引（范围：[-max_len+1, max_len-1]）
        position = torch.arange(-max_len + 1, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

        # 计算相对位置编码
        pe = torch.zeros(2 * max_len - 1, d_model)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        # 注册为buffer，不会被优化器更新
        self.register_buffer('pe', pe)

    def forward(self, qlen, klen):
        """
        qlen: 查询序列长度
        klen: 关键序列长度
        返回一个（qlen, klen, d_model）形状的相对位置编码
        """
        assert qlen <= self.max_len and self.max_len, "序列长度不能超过max_len"

        # 获取相对位置索引
        relative_position = torch.arange(klen).view(1, -1) - torch.arange(qlen).view(-1, 1)
        relative_position += self.max_len - 1  # 将索引转换为正数范围
        return self.pe[relative_position]  # (qlen, klen, d_model)


class RotaryPositionalEncoding(nn.Module):
    """旋转位置编码：核心思想是通过旋转操作在查询（Q）和键（K）向量中编码位置信息，而不直接对输入进行加性位置编码"""
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        self.d_model = d_model

        # 计算RoPE需要的theta(基于10000)
        position = torch.arange(max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        theta = position * div_term  # (max_len, d_model/2)

        # 计算sin和cos
        self.register_buffer("sin", torch.sin(theta))
        self.register_buffer("cos", torch.cos(theta))

    def forward(self, x):
        """
        x: 形状（batch, seq_len, d_model）
        """
        batch_size, seq_len, d_model = x.size()
        assert d_model == self.d_model, "输入的隐藏维度必须等于初始化时的d_model"

        # 仅对偶数维度应用旋转F
        x_even, x_odd = x[..., 0::2], x[..., 1::2]
        sin, cos = self.sin[:seq_len], self.cos[:seq_len]  # 取前seq_len长度的sin/cos

        # 旋转公式：
        x_rotated_even = x_even * cos - x_odd * sin
        x_rotated_odd = x_even * sin + x_odd * cos

        # 重新拼接回（batch, seq_len, d_model）
        x_rotated = torch.stack([x_rotated_even, x_rotated_odd], dim=-1).reshape(batch_size, seq_len, d_model)
        return x_rotated


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        # d_model：输入数据的特征维度，num_heads：多头注意力的头数，d_head：每个注意力头的特征维度
        self.d_head = d_model // num_heads
        self.num_heads = num_heads
        # W_q, W_k, W_v：三个线性变换层，用于将输入x投影到Q、K、V空间
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        # 输出层，负责把多头注意力的输出合并回d_model维度
        self.W_o = nn.Linear(d_model, d_model)
        # 丢弃层，用于防止过拟合
        self.dropout = nn.Dropout(0.1)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.d_head).transpose(1, 2)
        K = self.W_k(key).view(batch_size, -1, self.num_heads, self.d_head).transpose(1, 2)
        V = self.W_v(value).view(batch_size, -1, self.num_heads, self.d_head).transpose(1, 2)
        # Q、K、V都是经过线性变换后的输入张量，并按照多头划分
        # torch.matmul(Q, K.transpose(-1, -2))计算点积注意力的分数矩阵（scores）
        # math.sqrt(self.d_head)实现缩放，以避免梯度消失或梯度爆炸的问题
        scores = torch.matmul(Q, K.transpose(-1, -2)) / math.sqrt(self.d_head)
        # mask == 0的位置被填充为-inf，确保softmax计算时它们的概率变为0
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        # softmax归一化，使得每个位置的权重之和为1
        attention = torch.softmax(scores, dim=-1)
        attention = self.dropout(attention)
        # 加权求和并恢复形状
        # transpose(1, 2)交换num_heads和seq_len维度，变为(batch_size, seq_len, num_heads, d_head)
        # contiguous().view(batch_size, -1, d_model)变回(batch_size, seq_len, d_model)
        output = torch.matmul(attention, V).transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_head)
        # W_o负责融合多头信息，得到最终输出
        return self.W_o(output)


class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        # 第一层全连接，把d_model维的输入投影到d_ff维度的高维空间，提高特征表示能力
        self.fc1 = nn.Linear(d_model, d_ff)
        # 第二层全连接，把d_ff维的输入投影回d_model维度，恢复输入维度
        self.fc2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()
        self.gelu = nn.GELU()
        # 0.1的概率随机丢弃部分神经元，防止过拟合
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        """经过以下步骤
        1、self.fc1(x)：输入 x 经过第一个线性层，投影到 d_ff 维度。
        2、self.relu(...)：通过 ReLU 激活函数，引入非线性。
        3、self.dropout(...)：对激活后的结果进行 Dropout，防止过拟合。
        4、self.fc2(...)：再通过一个线性层，将 d_ff 维度的结果投影回 d_model 维度。
        """
        return self.fc2(self.dropout(self.gelu(self.fc1(x))))


class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask):
        # 自注意力层的残差连接
        # self.self_attn(x, x, x, mask)计算自注意力的输出
        # self.dropout(...)防止过拟合
        # x + ... 是残差连接，表示把输入x和计算得到的self_attn(x, x, x, mask)相加
        # self.norm1(...)进行层归一化（LayerNorm），防止分布漂移
        x = self.norm1(x + self.dropout(self.self_attn(x, x, x, mask)))
        # 前馈网络的残差连接
        # self.feed_forward(x)是两层的前馈网络（FFN）
        # self.dropout(...) 防止过拟合
        # x + ... 这里的x是自注意力层后的输出，它和FFN的输出相加，形成第二次残差连接
        # self.norm2(...) 进行层归一化，防止分布漂移
        x = self.norm2(x + self.dropout(self.feed_forward(x)))
        return x


class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = FeedForward(d_model, d_ff, dropout)
        # 用于Self-Attention之后
        self.norm1 = nn.LayerNorm(d_model)
        # 用于Cross-Attention之后
        self.norm2 = nn.LayerNorm(d_model)
        # 用于FFN之后
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, encoder_output, src_mask, tgt_mask):
        # 执行Masked注意力
        # self.self_attn(x, x, x, tgt_mask)：
        # 计算解码器的Masked 自注意力，其中 Q, K, V 都是 x，并使用 tgt_mask 进行掩码，防止看到未来的信息。
        # self.dropout(...)：
        # 对注意力结果进行 dropout 以增强正则化效果。
        # x + ...：
        # 残差连接（Residual Connection）：原始输入 x 加上计算结果，以便于梯度传播。
        # self.norm1(...)：
        # 层归一化（LayerNorm）：稳定训练并防止梯度消失或梯度爆炸。
        x = self.norm1(x + self.dropout(self.self_attn(x, x, x, tgt_mask)))
        # 执行交叉注意力
        # self.cross_attn(x, encoder_output, encoder_output, src_mask)：
        # 计算解码器对编码器输出的交叉注意力，其中：
        # Q 来自 x（当前解码序列）。
        # K, V 来自 encoder_output（编码器输出）。
        # src_mask 用于屏蔽编码器中的无效 token（如 padding）。
        # 残差连接 + 层归一化：与上一步类似。
        x = self.norm2(x + self.dropout(self.cross_attn(x, encoder_output, encoder_output, src_mask)))
        # 执行前馈神经网络
        # self.feed_forward(x)：
        # 计算 FFN 结果，将 d_model 映射到更高维度 d_ff，然后再降维回 d_model。
        # 残差连接 + 层归一化：与前两步相同。
        x = self.norm3(x + self.dropout(self.feed_forward(x)))
        return x


class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, num_heads=8, num_layers=6, d_ff=2048, dropout=0.1):
        super().__init__()
        self.encoder_embed = nn.Embedding(src_vocab_size, d_model)
        self.decoder_embed = nn.Embedding(tgt_vocab_size, d_model)
        self.pos_encoding = AbsolutePositionalEncoding(d_model)
        self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
        self.fc_out = nn.Linear(d_model, tgt_vocab_size)
        self.dropout = nn.Dropout(dropout)

    def encode(self, src, src_mask):
        src = self.dropout(self.pos_encoding(self.encoder_embed(src)))
        for layer in self.encoder_layers:
            src = layer(src, src_mask)
        return src

    def decode(self, tgt, encoder_output, src_mask, tgt_mask):
        tgt = self.dropout(self.pos_encoding(self.decoder_embed(tgt)))
        for layer in self.decoder_layers:
            tgt = layer(tgt, encoder_output, src_mask, tgt_mask)
        return self.fc_out(tgt)

    def forward(self, src, tgt, src_mask, tgt_mask):
        encoder_output = self.encode(src, src_mask)
        return self.decode(tgt, encoder_output, src_mask, tgt_mask)

    @torch.no_grad()
    # 前向传播时不构建计算图，节省内存并加速推理
    def generate(self, src, tgt_vocab, max_len=50):
        batch_size = src.size(0)
        device = src.device
        tgt = torch.ones(batch_size, 1).long().to(device)  # <sos> 索引为 1
        eos_idx = tgt_vocab.word2idx['<eos>']

        self.eval()  # 切换为评估模式，关闭Dropout/LayerNorm等操作
        src_mask = (src != 0).unsqueeze(1).unsqueeze(2)  # 只构造一次

        for _ in range(max_len):
            tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(2)
            output = self(src, tgt, src_mask, tgt_mask)
            next_word = output[:, -1, :].argmax(dim=-1)

            tgt = torch.cat([tgt, next_word.unsqueeze(1)], dim=1)

            if (next_word == eos_idx).all():
                break

        return tgt


# Debug: 检查 Transformer 是否正常运行
if __name__ == "__main__":
    model = Transformer(src_vocab_size=1000, tgt_vocab_size=1000, d_model=64, num_heads=2, num_layers=3, d_ff=256, dropout=0.1)
    src = torch.randint(0, 1000, (8, 64))  # batch_size=8, seq_len=64
    tgt = torch.randint(0, 1000, (8, 64))
    src_mask = torch.ones(8, 1, 1, 64)
    tgt_mask = torch.ones(8, 1, 64, 64)
    output = model(src, tgt, src_mask, tgt_mask)
    print(output.shape)  # 应该是 (8, 64, 1000)